#!/user/bin/env python3
# -*- coding: utf-8 -*-

import os

# 初始化Embedding类
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_core.vectorstores import InMemoryVectorStore
from langchain.document_loaders import TextLoader
from langchain_openai import ChatOpenAI

llm = ChatOpenAI(
    # 若没有配置环境变量，请用百炼API Key将下行替换为：api_key="sk-xxx",
    api_key=os.getenv("DASH_SCOPE_API_KEY"), # 如何获取API Key：https://help.aliyun.com/zh/model-studio/developer-reference/get-api-key
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    model="qwen-plus",
    temperature=0.8,
    max_tokens=60,
)

embeddings_model = DashScopeEmbeddings(
    model="text-embedding-v1",
    dashscope_api_key=os.getenv("DASH_SCOPE_API_KEY")
)

# 导入文档加载器模块，并使用TextLoader来加载文本文件
loader = TextLoader('../鲜花花语大全.txt', encoding='utf8')

# 使用VectorstoreIndexCreator来从加载器创建索引

# 导入内存存储库，该库允许我们在RAM中临时存储数据
from langchain.indexes import VectorstoreIndexCreator

index = VectorstoreIndexCreator(
    vectorstore_cls=InMemoryVectorStore, embedding=embeddings_model
).from_loaders([loader])

# 定义查询字符串, 使用创建的索引执行查询
query = "玫瑰花的花语是什么？"
result = index.query(llm=llm, question=query)
print(result) # 打印查询结果


# 替换成你所需要的工具
from langchain.text_splitter import CharacterTextSplitter
index_creator = VectorstoreIndexCreator(
    vectorstore_cls=InMemoryVectorStore,
    embedding=embeddings_model,
    text_splitter=CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
)
